import os
import time
import logging
import socket
from datetime import datetime, timedelta
from dotenv import load_dotenv
from tqdm import tqdm
import baostock as bs
import pymysql
from pymysql.cursors import DictCursor

# ------------------------------ 全局配置 ------------------------------
load_dotenv()

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('stock_k_combined_sync.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 数据库配置
DB_CONFIG = {
    'host': os.getenv('DB_HOST', 'localhost'),
    'user': os.getenv('DB_USER', 'root'),
    'password': os.getenv('DB_PASSWORD', '123456'),
    'database': os.getenv('DB_NAME', 'stock_data'),
    'port': int(os.getenv('DB_PORT', 3306)),
    'charset': 'utf8mb4'
}

# 股票市场收盘时间
MARKET_CLOSE_HOUR = 15
MARKET_CLOSE_MINUTE = 0

# 核心字段配置（修正周线/月线逻辑）
DAY_K_FIELDS = [
    'code', 'date', 'open', 'high', 'low', 'close',
    'volume', 'amount', 'adjustflag', 'turn', 'pctChg',
    'last_day_close'
]

BASE_K_FIELDS = [
    'code', 'date', 'open', 'high', 'low', 'close',
    'volume', 'amount', 'adjustflag', 'turn', 'pctChg'
]

K_CONFIG = {
    'day': {
        'table': 'stock_day_k_history',
        'progress_file': 'processed_day_stocks.txt',
        'freq': 'd',
        'fields': DAY_K_FIELDS,
        'last_period_close_field': 'last_day_close',
        'api_interval': 0.2,
        'days_limit': 200,
        'retry_count': 2
    },
    'week': {
        'table': 'stock_week_k_history',
        'progress_file': 'processed_week_stocks.txt',
        'freq': 'w',
        'fields': BASE_K_FIELDS + ['last_week_close'],
        'last_period_close_field': 'last_week_close',
        'api_interval': 0.3,
        'min_valid_weeks': 10,  # 接口返回周数少于此值则用日线计算
        'recent_threshold_days': 90  # 近3个月内有数据视为有效
    },
    'month': {
        'table': 'stock_month_k_history',
        'progress_file': 'processed_month_stocks.txt',
        'freq': 'm',
        'fields': BASE_K_FIELDS + ['last_month_close'],
        'last_period_close_field': 'last_month_close',
        'api_interval': 0.5,
        'min_valid_months': 3,  # 接口返回月数少于此值则用日线计算
        'recent_threshold_days': 180  # 近6个月内有数据视为有效
    }
}

# 重试配置
STOCK_LIST_RETRY = 10
RETRY_DELAY_BASE = 2
FORCE_REPROCESS = False
START_DATE = '1990-12-19'


# ------------------------------ 公共工具函数 ------------------------------
def init_combined_database():
    """初始化数据库表结构"""
    try:
        conn = pymysql.connect(
            host=DB_CONFIG['host'],
            user=DB_CONFIG['user'],
            password=DB_CONFIG['password'],
            port=DB_CONFIG['port'],
            charset=DB_CONFIG['charset']
        )
        cursor = conn.cursor()
        cursor.execute(f"CREATE DATABASE IF NOT EXISTS {DB_CONFIG['database']}")
        conn.select_db(DB_CONFIG['database'])

        for k_type in K_CONFIG.keys():
            cfg = K_CONFIG[k_type]
            if k_type == 'day':
                create_table_sql = f"""
                    CREATE TABLE IF NOT EXISTS {cfg['table']} (
                        id VARCHAR(50) NOT NULL COMMENT '唯一标识（股票代码_日期）',
                        code VARCHAR(20) NOT NULL COMMENT '股票代码（纯数字）',
                        date DATE NOT NULL COMMENT '日K线日期',
                        open DECIMAL(10, 2) NULL COMMENT '开盘价',
                        high DECIMAL(10, 2) NULL COMMENT '最高价',
                        low DECIMAL(10, 2) NULL COMMENT '最低价',
                        close DECIMAL(10, 2) NULL COMMENT '收盘价',
                        volume DECIMAL(12, 2) NULL COMMENT '成交量（万手）',
                        amount DECIMAL(12, 2) NULL COMMENT '成交额（亿元）',
                        adjustflag TINYINT NULL COMMENT '复权类型',
                        turn DECIMAL(10, 2) NULL COMMENT '换手率（%）',
                        pctChg DECIMAL(10, 2) NULL COMMENT '涨跌幅（%）',
                        last_day_close DECIMAL(10, 2) NULL COMMENT '上一日收盘价',
                        insert_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                        PRIMARY KEY (id),
                        INDEX idx_code_date (code, date),
                        INDEX idx_date (date)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
                """
            else:
                create_table_sql = f"""
                    CREATE TABLE IF NOT EXISTS {cfg['table']} (
                        id VARCHAR(50) NOT NULL COMMENT '唯一标识（股票代码_日期）',
                        code VARCHAR(20) NOT NULL COMMENT '股票代码（纯数字）',
                        date DATE NOT NULL COMMENT '{k_type}K线日期',
                        open DECIMAL(10, 2) NULL COMMENT '开盘价',
                        high DECIMAL(10, 2) NULL COMMENT '最高价',
                        low DECIMAL(10, 2) NULL COMMENT '最低价',
                        close DECIMAL(10, 2) NULL COMMENT '收盘价',
                        {cfg['last_period_close_field']} DECIMAL(10, 2) NULL COMMENT '上{cfg['table'].split('_')[1]}收盘价',
                        volume DECIMAL(12, 2) NULL COMMENT '成交量（万手）',
                        amount DECIMAL(12, 2) NULL COMMENT '成交额（亿元）',
                        adjustflag TINYINT NULL COMMENT '复权类型',
                        turn DECIMAL(10, 2) NULL COMMENT '换手率（%）',
                        pctChg DECIMAL(10, 2) NULL COMMENT '涨跌幅（%）',
                        insert_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                        PRIMARY KEY (id),
                        INDEX idx_code_date (code, date),
                        INDEX idx_date (date)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
                """
            cursor.execute(create_table_sql)
            logger.info(f"{k_type.upper()}K表初始化成功（表名：{cfg['table']}）")

        conn.commit()
        logger.info("=== 日K+周K+月K数据库表结构初始化完成 ===")
    except Exception as e:
        logger.error(f"数据库初始化失败：{str(e)}", exc_info=True)
        raise
    finally:
        if 'conn' in locals():
            conn.close()


def get_recent_trading_date(force_prev_day=False):
    """获取最近的交易日"""
    now = datetime.now()
    is_before_close = (now.hour < MARKET_CLOSE_HOUR) or \
                      (now.hour == MARKET_CLOSE_HOUR and now.minute < MARKET_CLOSE_MINUTE)

    start_offset = 1 if (is_before_close or force_prev_day) else 0
    logger.info(f"当前时间：{now.strftime('%Y-%m-%d %H:%M:%S')}，{'未收盘' if is_before_close else '已收盘'}，"
                f"{'从昨天开始' if start_offset == 1 else '从今天开始'}检查交易日")

    for i in range(start_offset, 30):
        check_date = (now - timedelta(days=i)).strftime('%Y-%m-%d')
        try:
            lg = bs.login()
            if lg.error_code != '0':
                logger.warning(f"查询交易日[{check_date}]登录失败：{lg.error_msg}，将重试")
                bs.logout()
                continue

            rs = bs.query_trade_dates(start_date=check_date, end_date=check_date)
            if rs.error_code != '0':
                logger.warning(f"查询交易日[{check_date}]失败：{rs.error_msg}")
                bs.logout()
                continue

            if rs.next():
                is_trading = rs.get_row_data()[1] == '1'
                bs.logout()
                if is_trading:
                    logger.info(f"找到有效交易日：{check_date}")
                    return check_date

            bs.logout()
            logger.debug(f"{check_date}不是交易日，继续检查前一天")
        except socket.error as se:
            logger.warning(f"查询交易日[{check_date}]套接字错误：{str(se)}，将重试")
            time.sleep(1)
        except Exception as e:
            logger.warning(f"查询交易日[{check_date}]异常：{str(e)}，将重试")
            time.sleep(1)

    safe_date = "2024-12-31"
    logger.warning(f"无法获取有效交易日，使用安全日期：{safe_date}")
    return safe_date


def get_all_stock_codes():
    """获取所有A股代码"""
    trade_date = get_recent_trading_date()

    for retry in range(STOCK_LIST_RETRY):
        if retry > 0:
            trade_date = get_recent_trading_date(force_prev_day=True)
            logger.info(f"第{retry + 1}次重试，使用更早的交易日：{trade_date}")

        retry_delay = RETRY_DELAY_BASE ** min(retry, 5)
        try:
            lg = bs.login()
            if lg.error_code != '0':
                logger.error(f"股票列表获取失败（登录错误，重试{retry + 1}/{STOCK_LIST_RETRY}）：{lg.error_msg}")
                bs.logout()
                time.sleep(retry_delay)
                continue

            socket.setdefaulttimeout(15)
            rs = bs.query_all_stock(trade_date)
            if rs.error_code != '0':
                logger.error(f"股票列表获取失败（接口错误，重试{retry + 1}/{STOCK_LIST_RETRY}）：{rs.error_msg}")
                bs.logout()
                time.sleep(retry_delay)
                continue

            stock_codes = []
            while (rs.error_code == '0') and rs.next():
                code = rs.get_row_data()[0]
                if code.startswith(('sh.60', 'sh.688', 'sz.00', 'sz.30')):
                    stock_codes.append(code)

            if len(stock_codes) < 3000:
                logger.warning(f"第{retry + 1}次获取到股票数量不足（{len(stock_codes)}只），可能数据异常，将重试")
                bs.logout()
                time.sleep(retry_delay)
                continue

            bs.logout()
            logger.info(f"=== 成功获取{len(stock_codes)}只A股股票代码（重试{retry + 1}/{STOCK_LIST_RETRY}） ===")
            return stock_codes

        except socket.timeout:
            logger.warning(f"股票列表获取超时（重试{retry + 1}/{STOCK_LIST_RETRY}），将重试")
            bs.logout()
            time.sleep(retry_delay)
        except socket.error as se:
            logger.warning(f"股票列表获取套接字错误（重试{retry + 1}/{STOCK_LIST_RETRY}）：{str(se)}，将重试")
            bs.logout()
            time.sleep(retry_delay)
        except Exception as e:
            logger.error(f"股票列表获取异常（重试{retry + 1}/{STOCK_LIST_RETRY}）：{str(e)}", exc_info=True)
            bs.logout()
            time.sleep(retry_delay)

    logger.error("=" * 50)
    logger.error("严重错误：所有尝试均未能从Baostock获取有效股票列表")
    logger.error("可能原因：")
    logger.error("1. Baostock服务临时故障（请稍后再试）")
    logger.error("2. 网络连接不稳定或被限制访问")
    logger.error("3. 近期无有效交易日（如长假期间）")
    logger.error("=" * 50)
    raise RuntimeError("无法获取股票列表，任务终止")


def load_valid_processed_stocks(k_type, all_stocks):
    """加载已处理股票"""
    cfg = K_CONFIG[k_type]
    if not os.path.exists(cfg['progress_file']):
        logger.info(f"{k_type.upper()}K进度文件不存在，将处理所有股票")
        return []

    with open(cfg['progress_file'], 'r', encoding='utf-8') as f:
        progress_codes = [line.strip() for line in f if line.strip()]

    valid_progress = [code for code in progress_codes if code in all_stocks]
    if not valid_progress:
        logger.info(f"{k_type.upper()}K进度文件中无有效股票代码，将处理所有股票")
        return []

    try:
        conn = pymysql.connect(**DB_CONFIG)
        cursor = conn.cursor()
        clean_codes = [code.split('.')[1] for code in valid_progress]
        placeholders = ', '.join(['%s'] * len(clean_codes))
        cursor.execute(f"""
            SELECT DISTINCT code FROM {cfg['table']} 
            WHERE code IN ({placeholders})
        """, tuple(clean_codes))
        db_codes = [row[0] for row in cursor.fetchall()]
        db_valid_codes = [code for code in valid_progress
                          if clean_stock_code(code) in db_codes]

        logger.info(f"{k_type.upper()}K验证后，已处理股票：{len(db_valid_codes)}只")
        return db_valid_codes
    except Exception as e:
        logger.warning(f"{k_type.upper()}K验证已处理股票失败，将重新处理：{str(e)}")
        return []
    finally:
        if 'conn' in locals():
            conn.close()


def save_processed_stock(k_type, code):
    """保存已处理股票"""
    cfg = K_CONFIG[k_type]
    with open(cfg['progress_file'], 'a', encoding='utf-8') as f:
        f.write(f"{code}\n")


def clean_stock_code(raw_code):
    """清洗股票代码"""
    return raw_code.split('.')[1] if '.' in raw_code else raw_code


def generate_id(code, date):
    """生成唯一ID"""
    date_str = date.replace('-', '')
    return f"{code}_{date_str}"


def convert_volume(volume):
    """转换成交量单位（股→万手）"""
    if volume is None:
        return None
    try:
        return round(volume / 1_000_000, 2)  # 1万手 = 100万股
    except (TypeError, ValueError):
        return None


def convert_amount(amount):
    """转换成交额单位（元→亿元）"""
    if amount is None:
        return None
    try:
        return round(amount / 100_000_000, 2)  # 1亿元 = 100000000元
    except (TypeError, ValueError):
        return None


def calculate_last_period_close(k_type, k_data):
    """计算上周期收盘价"""
    if not k_data:
        return []

    # 按日期排序
    sorted_data = sorted(k_data, key=lambda x: x['date'])
    cfg = K_CONFIG[k_type]
    last_close_field = cfg['last_period_close_field']

    for i in range(1, len(sorted_data)):
        prev_date = datetime.strptime(sorted_data[i - 1]['date'], '%Y-%m-%d')
        curr_date = datetime.strptime(sorted_data[i]['date'], '%Y-%m-%d')
        if curr_date > prev_date:  # 确保时间递增
            sorted_data[i][last_close_field] = sorted_data[i - 1]['close']
            # 计算涨跌幅（如果有上周期收盘价）
            if sorted_data[i][last_close_field] is not None and sorted_data[i][last_close_field] != 0:
                sorted_data[i]['pctChg'] = round(
                    (sorted_data[i]['close'] - sorted_data[i][last_close_field]) /
                    sorted_data[i][last_close_field] * 100, 2
                )
        else:
            logger.warning(
                f"{k_type.upper()}K数据乱序：{sorted_data[i]['date']} ≤ {sorted_data[i - 1]['date']}，跳过计算")
            sorted_data[i][last_close_field] = None

    # 第一条数据无前期收盘价
    sorted_data[0][last_close_field] = None
    sorted_data[0]['pctChg'] = None
    return sorted_data


# ------------------------------ K线计算逻辑（新增核心功能） ------------------------------
def calculate_week_k_from_day_k(day_k_data):
    """用日线数据计算周线（按自然周划分）"""
    if not day_k_data:
        return []

    # 按周分组（key: 年份+周数，value: 当周日线数据）
    week_groups = {}
    for day in day_k_data:
        date = datetime.strptime(day['date'], '%Y-%m-%d')
        year, week_num, _ = date.isocalendar()  # 年份+周数作为分组key
        key = f"{year}-W{week_num:02d}"
        if key not in week_groups:
            week_groups[key] = []
        week_groups[key].append(day)

    week_k_list = []
    for key, days in week_groups.items():
        # 按日期排序，取每周最后一个交易日作为周线日期
        sorted_days = sorted(days, key=lambda x: x['date'])
        week_end_date = sorted_days[-1]['date']  # 周线日期为周末最后一个交易日
        code = days[0]['code']

        # 计算周线指标
        week_k = {
            'id': generate_id(code, week_end_date),
            'code': code,
            'date': week_end_date,
            'open': sorted_days[0]['open'],  # 周开盘价=周内首个交易日开盘价
            'high': max(day['high'] for day in sorted_days if day['high'] is not None),  # 周最高价
            'low': min(day['low'] for day in sorted_days if day['low'] is not None),  # 周最低价
            'close': sorted_days[-1]['close'],  # 周收盘价=周内最后一个交易日收盘价
            'volume': sum(day['volume'] for day in sorted_days if day['volume'] is not None),  # 周成交量
            'amount': sum(day['amount'] for day in sorted_days if day['amount'] is not None),  # 周成交额
            'adjustflag': sorted_days[0]['adjustflag'],  # 复权类型沿用日线
            'turn': sum(day['turn'] for day in sorted_days if day['turn'] is not None),  # 周换手率
            'pctChg': None  # 涨跌幅后续计算
        }
        week_k_list.append(week_k)

    # 计算上周收盘价和涨跌幅
    return calculate_last_period_close('week', week_k_list)


def calculate_month_k_from_day_k(day_k_data):
    """用日线数据计算月线（按自然月划分）"""
    if not day_k_data:
        return []

    # 按月分组（key: 年月，value: 当月日线数据）
    month_groups = {}
    for day in day_k_data:
        date = datetime.strptime(day['date'], '%Y-%m-%d')
        key = f"{date.year}-{date.month:02d}"  # 年月作为分组key
        if key not in month_groups:
            month_groups[key] = []
        month_groups[key].append(day)

    month_k_list = []
    for key, days in month_groups.items():
        # 按日期排序，取每月最后一个交易日作为月线日期
        sorted_days = sorted(days, key=lambda x: x['date'])
        month_end_date = sorted_days[-1]['date']  # 月线日期为月末最后一个交易日
        code = days[0]['code']

        # 计算月线指标
        month_k = {
            'id': generate_id(code, month_end_date),
            'code': code,
            'date': month_end_date,
            'open': sorted_days[0]['open'],  # 月开盘价=月内首个交易日开盘价
            'high': max(day['high'] for day in sorted_days if day['high'] is not None),  # 月最高价
            'low': min(day['low'] for day in sorted_days if day['low'] is not None),  # 月最低价
            'close': sorted_days[-1]['close'],  # 月收盘价=月内最后一个交易日收盘价
            'volume': sum(day['volume'] for day in sorted_days if day['volume'] is not None),  # 月成交量
            'amount': sum(day['amount'] for day in sorted_days if day['amount'] is not None),  # 月成交额
            'adjustflag': sorted_days[0]['adjustflag'],  # 复权类型沿用日线
            'turn': sum(day['turn'] for day in sorted_days if day['turn'] is not None),  # 月换手率
            'pctChg': None  # 涨跌幅后续计算
        }
        month_k_list.append(month_k)

    # 计算上月收盘价和涨跌幅
    return calculate_last_period_close('month', month_k_list)


# ------------------------------ 核心同步函数 ------------------------------
def _fetch_week_k_from_api(raw_code, cfg, global_login):
    """从接口获取周线数据（内部封装）"""
    try:
        start_date = START_DATE
        fields_str = ",".join(BASE_K_FIELDS)

        rs = bs.query_history_k_data_plus(
            raw_code,
            fields_str,
            start_date=start_date,
            end_date=datetime.now().strftime('%Y-%m-%d'),
            frequency=cfg['freq'],
            adjustflag='2'  # 后复权
        )

        if rs.error_code != '0':
            logger.warning(f"WEEKK：股票{raw_code}接口错误：{rs.error_msg}")
            return []

        # 解析接口数据
        data_list = []
        clean_code = clean_stock_code(raw_code)
        while (rs.error_code == '0') and rs.next():
            row = rs.get_row_data()
            raw_data = dict(zip(BASE_K_FIELDS, row))

            processed_data = {
                'id': generate_id(clean_code, raw_data['date']),
                'code': clean_code,
                'date': raw_data['date']
            }

            # 基础字段处理
            for key in ['open', 'high', 'low', 'close', 'turn', 'pctChg']:
                processed_data[key] = float(raw_data[key]) if raw_data.get(key) and raw_data[key] != '' else None

            processed_data['volume'] = convert_volume(
                int(raw_data['volume']) if raw_data.get('volume') and raw_data['volume'] != '' else None
            )
            processed_data['amount'] = convert_amount(
                float(raw_data['amount']) if raw_data.get('amount') and raw_data['amount'] != '' else None
            )
            processed_data['adjustflag'] = int(raw_data['adjustflag']) if raw_data.get('adjustflag') else None

            data_list.append(processed_data)

        # 计算上周收盘价
        return calculate_last_period_close('week', data_list)
    except Exception as e:
        logger.error(f"WEEKK：股票{raw_code}接口数据解析失败：{str(e)}")
        return []


def _fetch_month_k_from_api(raw_code, cfg, global_login):
    """从接口获取月线数据（内部封装）"""
    try:
        start_date = START_DATE
        fields_str = ",".join(BASE_K_FIELDS)

        rs = bs.query_history_k_data_plus(
            raw_code,
            fields_str,
            start_date=start_date,
            end_date=datetime.now().strftime('%Y-%m-%d'),
            frequency=cfg['freq'],
            adjustflag='2'  # 后复权
        )

        if rs.error_code != '0':
            logger.warning(f"MONTMK：股票{raw_code}接口错误：{rs.error_msg}")
            return []

        # 解析接口数据
        data_list = []
        clean_code = clean_stock_code(raw_code)
        while (rs.error_code == '0') and rs.next():
            row = rs.get_row_data()
            raw_data = dict(zip(BASE_K_FIELDS, row))

            processed_data = {
                'id': generate_id(clean_code, raw_data['date']),
                'code': clean_code,
                'date': raw_data['date']
            }

            # 基础字段处理
            for key in ['open', 'high', 'low', 'close', 'turn', 'pctChg']:
                processed_data[key] = float(raw_data[key]) if raw_data.get(key) and raw_data[key] != '' else None

            processed_data['volume'] = convert_volume(
                int(raw_data['volume']) if raw_data.get('volume') and raw_data['volume'] != '' else None
            )
            processed_data['amount'] = convert_amount(
                float(raw_data['amount']) if raw_data.get('amount') and raw_data['amount'] != '' else None
            )
            processed_data['adjustflag'] = int(raw_data['adjustflag']) if raw_data.get('adjustflag') else None

            data_list.append(processed_data)

        # 计算上月收盘价
        return calculate_last_period_close('month', data_list)
    except Exception as e:
        logger.error(f"MONTMK：股票{raw_code}接口数据解析失败：{str(e)}")
        return []


def _has_recent_data(k_data, threshold_days):
    """检查是否包含近期数据"""
    if not k_data:
        return False
    latest_date = datetime.strptime(k_data[-1]['date'], '%Y-%m-%d')
    return (datetime.now() - latest_date) <= timedelta(days=threshold_days)


def fetch_stock_k_data(k_type, raw_code, global_login=False):
    """获取单只股票K线数据（周线/月线支持接口+日线计算兜底）"""
    cfg = K_CONFIG[k_type]
    try:
        if not global_login:
            lg = bs.login()
            if lg.error_code != '0':
                logger.warning(f"{k_type.upper()}K：股票{raw_code}登录失败，跳过：{lg.error_msg}")
                return []

        # 周线逻辑：接口优先，不足则用日线计算
        if k_type == 'week':
            # 1. 尝试接口获取
            week_data = _fetch_week_k_from_api(raw_code, cfg, global_login)
            # 检查接口数据有效性
            if len(week_data) >= cfg['min_valid_weeks'] and \
                    _has_recent_data(week_data, cfg['recent_threshold_days']):
                logger.info(f"WEEKK：股票{raw_code}使用接口数据（{len(week_data)}条）")
                return week_data
            # 2. 接口数据无效，用日线计算
            logger.warning(f"WEEKK：股票{raw_code}接口数据不足（{len(week_data)}条），改用日线计算")
            day_data = fetch_stock_k_data('day', raw_code, global_login)  # 复用日线数据
            return calculate_week_k_from_day_k(day_data)

        # 月线逻辑：接口优先，不足则用日线计算
        if k_type == 'month':
            # 1. 尝试接口获取
            month_data = _fetch_month_k_from_api(raw_code, cfg, global_login)
            # 检查接口数据有效性
            if len(month_data) >= cfg['min_valid_months'] and \
                    _has_recent_data(month_data, cfg['recent_threshold_days']):
                logger.info(f"MONTMK：股票{raw_code}使用接口数据（{len(month_data)}条）")
                return month_data
            # 2. 接口数据无效，用日线计算
            logger.warning(f"MONTMK：股票{raw_code}接口数据不足（{len(month_data)}条），改用日线计算")
            day_data = fetch_stock_k_data('day', raw_code, global_login)  # 复用日线数据
            return calculate_month_k_from_day_k(day_data)

        # 日线逻辑（保持不变）
        if k_type == 'day':
            start_date = (datetime.now() - timedelta(days=cfg['days_limit'])).strftime('%Y-%m-%d')
            fields_str = ",".join([f for f in cfg['fields'] if f != cfg['last_period_close_field']])

            rs = bs.query_history_k_data_plus(
                raw_code,
                fields_str,
                start_date=start_date,
                end_date=datetime.now().strftime('%Y-%m-%d'),
                frequency=cfg['freq'],
                adjustflag='2'  # 后复权
            )

            if rs.error_code != '0':
                logger.warning(f"DAYK：股票{raw_code}数据获取失败：{rs.error_msg}")
                return []

            # 解析日线数据
            data_list = []
            clean_code = clean_stock_code(raw_code)
            while (rs.error_code == '0') and rs.next():
                row = rs.get_row_data()
                raw_fields = [f for f in cfg['fields'] if f != cfg['last_period_close_field']]
                raw_data = dict(zip(raw_fields, row))

                processed_data = {
                    'id': generate_id(clean_code, raw_data['date']),
                    'code': clean_code,
                    'date': raw_data['date']
                }

                # 基础字段处理
                for key in ['open', 'high', 'low', 'close', 'turn', 'pctChg']:
                    processed_data[key] = float(raw_data[key]) if raw_data.get(key) and raw_data[key] != '' else None

                processed_data['volume'] = convert_volume(
                    int(raw_data['volume']) if raw_data.get('volume') and raw_data['volume'] != '' else None
                )
                processed_data['amount'] = convert_amount(
                    float(raw_data['amount']) if raw_data.get('amount') and raw_data['amount'] != '' else None
                )
                processed_data['adjustflag'] = int(raw_data['adjustflag']) if raw_data.get('adjustflag') else None

                data_list.append(processed_data)

            # 计算上日收盘价
            return calculate_last_period_close('day', data_list)

    except Exception as e:
        logger.error(f"{k_type.upper()}K：股票{raw_code}处理异常，跳过：{str(e)}", exc_info=True)
        return []
    finally:
        if not global_login:
            bs.logout()


def batch_insert_to_db(k_type, data_list):
    """批量插入/更新数据库（支持更新已有数据）"""
    if not data_list:
        return 0

    cfg = K_CONFIG[k_type]
    try:
        conn = pymysql.connect(**DB_CONFIG, cursorclass=DictCursor)
        cursor = conn.cursor()
        # 构建插入字段
        insert_fields = ['id', 'code', 'date'] + [f for f in cfg['fields'] if f not in ['id', 'code', 'date']]

        # 构建更新语句（主键冲突时更新所有字段）
        update_fields = [f"{field}=VALUES({field})" for field in insert_fields if field != 'id']
        sql = f"""
            INSERT INTO {cfg['table']} ({','.join(insert_fields)})
            VALUES ({','.join(['%s'] * len(insert_fields))})
            ON DUPLICATE KEY UPDATE {','.join(update_fields)}
        """

        # 批量处理
        batch_size = 2000
        total_affected = 0
        for i in range(0, len(data_list), batch_size):
            batch = data_list[i:i + batch_size]
            values = [tuple(item[field] for field in insert_fields) for item in batch]
            affected = cursor.executemany(sql, values)
            total_affected += affected
            conn.commit()

        logger.info(f"{k_type.upper()}K：批量操作完成，影响{total_affected}条数据（新增+更新）")
        return total_affected
    except Exception as e:
        conn.rollback()
        logger.error(f"{k_type.upper()}K：数据库操作失败：{str(e)}", exc_info=True)
        return 0
    finally:
        if 'conn' in locals():
            conn.close()


def sync_single_k_type(k_type, all_stocks, global_login_success):
    """单类型K线同步"""
    cfg = K_CONFIG[k_type]
    logger.info(f"\n===== 开始{cfg['freq'].upper()}K线同步任务 =====")
    start_time = time.time()

    if FORCE_REPROCESS:
        need_process = all_stocks
    else:
        processed_stocks = load_valid_processed_stocks(k_type, all_stocks)
        need_process = [code for code in all_stocks if code not in processed_stocks]

    logger.info(f"{k_type.upper()}K：股票总数：{len(all_stocks)}只")
    logger.info(f"{k_type.upper()}K：需处理股票：{len(need_process)}只")
    if not need_process and not FORCE_REPROCESS:
        logger.info(f"{k_type.upper()}K：所有股票均已同步完成，跳过任务")
        return 0, 0

    total_success = 0
    total_data = 0

    for raw_code in tqdm(need_process, desc=f"{k_type.upper()}K处理进度"):
        time.sleep(cfg['api_interval'])
        clean_code = clean_stock_code(raw_code)

        # 多轮重试机制
        k_data = None
        for retry in range(cfg.get('retry_count', 1) + 1):
            k_data = fetch_stock_k_data(k_type, raw_code, global_login_success)
            if len(k_data) > 0 or retry >= cfg.get('retry_count', 1):
                break
            logger.warning(f"{k_type.upper()}K：股票{raw_code}第{retry + 1}次重试（当前数据量：0）")
            time.sleep(1)  # 重试间隔

        if not k_data:
            logger.warning(
                f"{k_type.upper()}K：股票{raw_code}经{cfg.get('retry_count', 1) + 1}次尝试仍未获取到数据，跳过")
            continue

        # 插入数据库（支持更新）
        affected = batch_insert_to_db(k_type, k_data)
        if affected <= 0:
            logger.warning(f"{k_type.upper()}K：股票{raw_code}首次插入失败，重试1次")
            time.sleep(1)
            affected = batch_insert_to_db(k_type, k_data)
            if affected <= 0:
                logger.error(f"{k_type.upper()}K：股票{raw_code}多次插入失败，放弃")
                continue

        # 标记为已处理
        save_processed_stock(k_type, raw_code)
        total_success += 1
        total_data += len(k_data)
        logger.info(f"{k_type.upper()}K：股票{raw_code}处理完成（累计：{total_success}/{len(need_process)}）")

    end_time = time.time()
    elapsed = int(end_time - start_time)
    logger.info(f"""
    ===== {k_type.upper()}K线同步任务完成 =====
    总处理股票数：{len(need_process)}只
    成功同步股票数：{total_success}只
    总插入数据量：{total_data}条
    耗时：{elapsed // 3600}小时{(elapsed % 3600) // 60}分钟{elapsed % 60}秒
    """)
    return total_success, total_data


# ------------------------------ 主函数 ------------------------------
def main():
    logger.info("=" * 50)
    logger.info("===== 开始日K→周K→月K连续同步任务 =====")
    logger.info("=" * 50)
    overall_start = time.time()

    # 初始化数据库
    init_combined_database()

    # 获取股票列表
    all_stocks = get_all_stock_codes()
    if not all_stocks:
        logger.error("未获取到任何股票代码，任务终止")
        return

    # 全局登录（减少重复登录开销）
    global_login_success = False
    try:
        lg = bs.login()
        if lg.error_code == '0':
            global_login_success = True
            logger.info("\n=== Baostock全局登录成功 ===")
        else:
            logger.warning(f"\n=== Baostock全局登录失败，将使用单股票登录模式：{lg.error_msg} ===")

        # 依次同步日K→周K→月K（周/月K依赖日线数据）
        day_success, day_data = sync_single_k_type('day', all_stocks, global_login_success)
        week_success, week_data = sync_single_k_type('week', all_stocks, global_login_success)
        month_success, month_data = sync_single_k_type('month', all_stocks, global_login_success)

        # 任务总结
        overall_end = time.time()
        overall_elapsed = int(overall_end - overall_start)
        logger.info("\n" + "=" * 50)
        logger.info("===== 日K→周K→月K连续同步任务全部完成 =====")
        logger.info(f"日K同步：成功{day_success}只股票，插入{day_data}条数据")
        logger.info(f"周K同步：成功{week_success}只股票，插入{week_data}条数据")
        logger.info(f"月K同步：成功{month_success}只股票，插入{month_data}条数据")
        logger.info(f"总耗时：{overall_elapsed // 3600}小时{(overall_elapsed % 3600) // 60}分钟{overall_elapsed % 60}秒")
        logger.info("=" * 50)

    finally:
        if global_login_success:
            bs.logout()
            logger.info("\n=== Baostock全局登出成功 ===")


if __name__ == "__main__":
    main()